package 左哥算法.ch06二叉树;

import com.alibaba.fastjson.JSON;

import java.util.*;

/**
 * 来自于leetcode：1382,将二叉树变平衡
 */
public class AVLTree2 {
    Map<TreeNode,TreeNode> parentMap=new HashMap<>();
    Map<TreeNode,Integer> heightMap=new HashMap<>();
    TreeNode treeRoot=null;

    public TreeNode balanceBST(TreeNode root) {
        if (root==null){
            return null;
        }
        createTree(root);
        return treeRoot;
    }
    //以先序方式添加到树中
    public void createTree(TreeNode node){
        if (node==null){
            return;
        }
        insert(new TreeNode(node.val));
        createTree(node.left);
        createTree(node.right);
    }

    public void insert(TreeNode node){
        if (treeRoot==null){
            treeRoot=node;
            return;
        }
        if (node==null){
            return;
        }
        TreeNode curr=treeRoot;
        TreeNode parent=null;
        while (curr!=null){
            parent=curr;
            curr=curr.val>node.val?curr.left:curr.right;
        }
        if (parent.val>node.val){
            parent.left=node;
        }else {
            parent.right=node;
        }
        parentMap.put(node,parent);
        heightMap.put(node,1);
        this.treeRoot=maintain(parent);
    }

    /**
     * 使oldNode的父引用正确的指向新节点
     */
    public void replace(TreeNode oldNode,TreeNode newNode){
        TreeNode parent = parentMap.get(oldNode);
        if (oldNode==null|| parent ==null){
            return;
        }
        if (parent.left==oldNode){
            parent.left=newNode;
        }else {
            parent.right=newNode;
        }
    }

    public TreeNode maintain(TreeNode node){
        updateHeight(node);
        TreeNode curr=node;
        TreeNode newRoot=null;
        TreeNode pre=curr;
        while (curr!=null){
            int leftH=heightMap.getOrDefault(curr.left,0);
            int rightH=heightMap.getOrDefault(curr.right,0);
            if (leftH-rightH>1){
                int L=heightMap.getOrDefault(curr.left.left,0);
                int R=heightMap.getOrDefault(curr.left.right,0);
                if (L>R){   //LL型
                    curr = rightRotate(curr);
                    updateHeight(curr.right);
                }else{      //LR型
                    curr.left=leftRotate(curr.left);
                    curr= rightRotate(curr);
                    updateHeight(curr.left);
                    updateHeight(curr.right);
                }
            }else if (rightH-leftH>1){
                int R=heightMap.getOrDefault(curr.right.right,0);
                int L=heightMap.getOrDefault(curr.right.left,0);
                if (R>L){   //RR型
                    curr = leftRotate(curr);
                    updateHeight(curr.left);
                }else{      //RL型
                    curr.right=rightRotate(curr.right);
                    curr= leftRotate(curr);
                    updateHeight(curr.left);
                    updateHeight(curr.right);
                }
            }
            pre=curr;
            curr=parentMap.get(curr);
        }
        return pre;
    }

    public TreeNode leftRotate(TreeNode node){
        if (node==null||node.right==null){
            return node;
        }
        TreeNode root = node.right;
        replace(node,root);
        node.right=root.left;
        root.left=node;
        parentMap.put(node.right,node);
        parentMap.put(root,parentMap.get(node));
        parentMap.put(node,root);
        return root;
    }

    public TreeNode rightRotate(TreeNode node){
        if (node==null||node.left==null){
            return node;
        }
        TreeNode root = node.left;
        replace(node,root);
        node.left=root.right;
        root.right=node;
        parentMap.put(node.left,node);
        parentMap.put(root,parentMap.get(node));
        parentMap.put(node,root);
        return root;
    }


    public void updateHeight(TreeNode node){
        while (node != null) {
            int childHeight=Math.max(heightMap.getOrDefault(node.left,0),heightMap.getOrDefault(node.right,0));
            heightMap.put(node,childHeight+1);
            node=parentMap.get(node);
        }
    }

}

class Answer {

    public TreeNode balanceBST(TreeNode root) {
        if (root == null){
            return null;
        }
        // node节点的高度缓存
        Map<TreeNode,Integer> nodeHeight = new HashMap<>();
        TreeNode newRoot = null;
        Deque<TreeNode> stack = new LinkedList<>();
        TreeNode node = root;
        // 先序遍历插入（其实用哪个遍历都行）
        while(node != null || !stack.isEmpty()){
            if (node != null){
                // 新树插入
                newRoot = insert(newRoot,node.val,nodeHeight);
                stack.push(node);
                node = node.left;
            }else {
                node = stack.pop();
                node = node.right;
            }
        }
        return newRoot;
    }

    /**
     * 新节点插入
     * @param root root
     * @param val 新加入的值
     * @param nodeHeight 节点高度缓存
     * @return 新的root节点
     */
    private TreeNode insert(TreeNode root,int val,Map<TreeNode,Integer> nodeHeight){
        if (root == null){
            root = new TreeNode(val);
            nodeHeight.put(root,1);// 新节点的高度
            return root;
        }
        TreeNode node = root;
        int cmp = val - node.val;
        if (cmp < 0){
            // 左子树插入
            node.left = insert(root.left,val,nodeHeight);
            // 如果左右子树高度差超过1，进行旋转调整
            if (nodeHeight.getOrDefault(node.left,0) - nodeHeight.getOrDefault(node.right,0) > 1){
                if (val > node.left.val){
                    // 插入在左孩子右边，左孩子先左旋
                    node.left = rotateLeft(node.left,nodeHeight);
                }
                // 节点右旋
                node = rotateRight(node,nodeHeight);
            }
        }else if (cmp > 0){
            // 右子树插入
            node.right = insert(root.right,val,nodeHeight);
            // 如果左右子树高度差超过1，进行旋转调整
            if (nodeHeight.getOrDefault(node.right,0) - nodeHeight.getOrDefault(node.left,0) > 1){
                if (val < node.right.val){
                    // 插入在右孩子左边，右孩子先右旋
                    node.right = rotateRight(node.right,nodeHeight);
                }
                // 节点左旋
                node = rotateLeft(node,nodeHeight);
            }
        }else {
            // 一样的节点，啥都没发生
            return node;
        }
        // 获取当前节点新高度
        int height =  getCurNodeNewHeight(node,nodeHeight);
        // 更新当前节点高度
        nodeHeight.put(node,height);
        return node;
    }

    /**
     * node节点左旋
     * @param node node
     * @param nodeHeight node高度缓存
     * @return 旋转后的当前节点
     */
    private TreeNode rotateLeft(TreeNode node,Map<TreeNode,Integer> nodeHeight){
        // ---指针调整
        TreeNode right = node.right;
        node.right = right.left;
        right.left = node;
        // ---高度更新
        // 先更新node节点的高度，这个时候node是right节点的左孩子
        int newNodeHeight = getCurNodeNewHeight(node,nodeHeight);
        // 更新node节点高度
        nodeHeight.put(node,newNodeHeight);
        // newNodeHeight是现在right节点左子树高度，原理一样，取现在right左右子树最大高度+1
        int newRightHeight = Math.max(newNodeHeight,nodeHeight.getOrDefault(right.right,0)) + 1;
        // 更新原right节点高度
        nodeHeight.put(right,newRightHeight);
        return right;
    }

    /**
     * node节点右旋
     * @param node node
     * @param nodeHeight node高度缓存
     * @return 旋转后的当前节点
     */
    private TreeNode rotateRight(TreeNode node,Map<TreeNode,Integer> nodeHeight){
        // ---指针调整
        TreeNode left = node.left;
        node.left = left.right;
        left.right = node;
        // ---高度更新
        // 先更新node节点的高度，这个时候node是right节点的左孩子
        int newNodeHeight = getCurNodeNewHeight(node,nodeHeight);
        // 更新node节点高度
        nodeHeight.put(node,newNodeHeight);
        // newNodeHeight是现在left节点右子树高度，原理一样，取现在right左右子树最大高度+1
        int newLeftHeight = Math.max(newNodeHeight,nodeHeight.getOrDefault(left.left,0)) + 1;
        // 更新原left节点高度
        nodeHeight.put(left,newLeftHeight);
        return left;
    }

    /**
     * 获取当前节点的新高度
     * @param node node
     * @param nodeHeight node高度缓存
     * @return 当前node的新高度
     */
    private int getCurNodeNewHeight(TreeNode node,Map<TreeNode,Integer> nodeHeight){
        // node节点的高度，为现在node左右子树最大高度+1
        return Math.max(nodeHeight.getOrDefault(node.left,0),nodeHeight.getOrDefault(node.right,0)) + 1;
    }
}

class Test{
    public static void main(String[] args) {
        compare();
    }
    public static Set<Integer> set=new HashSet<>();
    static Random random=new Random();
    public static void compare(){
        int count=10000;
        while (count-- > 0) {
            int rest= random.nextInt(20)+1;  //节点层数
            int num=generateNum();
            TreeNode root = new TreeNode(num);
            createTree(root,rest);
            Answer answer=new Answer();
            AVLTree2 myFun=new AVLTree2();
            TreeNode root1 = myFun.balanceBST(root);
            TreeNode root2 = answer.balanceBST(root);
            if (!preOrder(root1,root2)){
                System.out.println("存在错误：");
                System.out.println(rest);
                System.out.println(JSON.toJSONString(root));
                List<Integer> list = preorder01(root);
                System.out.println(list.toString());
            }
            set.clear();
        }
    }
    public static List<Integer> preorder01(TreeNode root) {
        List<Integer> list = new ArrayList<>();
        if (root==null){
            return list;
        }
        list.add(root.val);
        list.addAll(preorder01(root.left));
        list.addAll(preorder01(root.right));
        return list;
    }

    public static boolean preOrder(TreeNode root1, TreeNode root2){
        if (root1==null&&root2==null){
            return true;
        }
        if (root1==null||root2==null){
            return false;
        }
        if (root1.val!=root2.val){
            return false;
        }
        boolean a = preOrder(root1.left, root2.left);
        boolean b = preOrder(root1.right, root2.right);
        return a&&b;
    }

    public static void createTree(TreeNode node,int rest){
        if (rest<=0||node==null){
            return;
        }
        node.left=random.nextBoolean()?new TreeNode(generateNum()):null;
        node.right=random.nextBoolean()?new TreeNode(generateNum()):null;
        createTree(node.left,rest-random.nextInt(2)-1);
        createTree(node.right,rest-random.nextInt(2)-1);
    }
    public static int generateNum(){
        int num;
        do {
            num=random.nextInt(200)+1;
        }while (set.contains(num));
        set.add(num);
        return num;
    }
}
